from unstructured.partition.pdf import partition_pdf
from unstructured.partition.docx import partition_docx
from unstructured.partition.doc import partition_doc
from unstructured.partition.ppt import partition_ppt
from unstructured.partition.text import partition_text
from unstructured.chunking.title import chunk_by_title

from chunks.BaseChunk import BaseChunk

SUPPORT_FILE_TYPE = {
    "pdf": partition_pdf,
    "doc": partition_doc,
    "docx": partition_docx,
    "ppt": partition_ppt,
    "text": partition_text
}


class UnstructuredChunk(BaseChunk):
    def __init__(self, max_characters: int = 500, overlap: int = 100):
        self.max_characters = min(2000, max_characters)
        self.overlap = min(1000, overlap)
        super().__init__()

    def getType(self):
        type = self.filename.split(".")[-1]
        if type not in SUPPORT_FILE_TYPE.keys():
            return None
        return type   
        
    def load(self, path: str):
        self.filepath = path
        self.filename = path.split("/")[-1]
        self.filetype = self.getType()
        if self.filetype == None:
            return False
        return True
        
    def chunk(self):
        partitionFunc = SUPPORT_FILE_TYPE[self.filetype]
        self.elements = partitionFunc(self.filepath)
        self.chunks = chunk_by_title(self.elements, max_characters=self.max_characters, \
            overlap = self.overlap)

        self.pageList = []
        self.chunkList = []
        
        for i in self.chunks:
            self.chunkList.append(i.text)
            self.pageList.append(i.metadata.orig_elements[0].metadata.page_number)
            
        return [self.chunkList, self.pageList]